# -*- coding: utf-8 -*-
from __future__ import division
from __future__ import print_function

from ..ops import Ref, Mean, Sum, Abs
from ..base import OpsFeature

from .close import Close
from .high import High
from .low import Low
from .tr import TR

class DT1(OpsFeature):
    """Technical Indicator - DMI (DT1)

    - HD = HIGH - REF(HIGH, 1)
    - LD = REF(LOW, 1) - LOW
    - DMP = SUM(IF(HD>0 AND HD>LD, HD, 0), N1)
    - DI1 = DMP / TR
    """
    def __init__(self, N=14):
        self.N = N

    def __str__(self):
        name = type(self).__name__.lower()
        if self.N != 14:
            name += '_'+str(self.N)
        return name

    def _load_internal(self, instrument, start_time, end_time, freq):
        tr = Mean(TR(), self.N)
        hd = High() - Ref(High(), 1)
        ld = Ref(Low(), 1) - Low()
        dmp = Sum(((hd > 0) & (hd > ld)) * hd, self.N)
        dt1 = dmp / tr
        return dt1.load(instrument, start_time, end_time, freq)
        

class DT2(OpsFeature):
    """Technical Indicator - DMI (DT2)

    - HD = HIGH - REF(HIGH, 1)
    - LD = REF(LOW, 1) - LOW
    - DMM = SUM(IF(LD>0 AND LD>HD, LD, 0), N1)
    - DI2 = DMM / TR
    """
    def __init__(self, N=14):
        self.N = N

    def __str__(self):
        name = type(self).__name__.lower()
        if self.N != 14:
            name += '_'+str(self.N)
        return name

    def _load_internal(self, instrument, start_time, end_time, freq):
        tr = Mean(TR(), self.N)
        hd = High() - Ref(High(), 1)
        ld = Ref(Low(), 1) - Low()
        dmm = Sum(((ld > 0) & (ld > hd)) * ld, self.N)
        dt2 = dmm / tr
        return dt2.load(instrument, start_time, end_time, freq)

    
class ADX(OpsFeature):
    """Technical Indicator - DMI (ADX)

    - ADX = MA(ABS(DI2 - DI1) / (DI1 + DI2), N2)
    """
    def __init__(self, N1=14, N2=6):
        self.N1 = N1
        self.N2 = N2

    def __str__(self):
        name = type(self).__name__.lower()
        if self.N1 != 14 or self.N2 != 6:
            name += '_{}_{}'.format(self.N1, self.N2)
        return name
    
    def _load_internal(self, instrument, start_time, end_time, freq):
        dt1 = DT1(self.N1)
        dt2 = DT2(self.N1)
        adx = Mean(Abs(dt2 - dt1) / (dt1 + dt2), self.N2)
        return adx.load(instrument, start_time, end_time, freq)


class ADXR(OpsFeature):
    """Technical Indicator - DMI (ADXR)

    - ADXR = (ADX + REF(ADX, N2)) / 2
    """
    def __init__(self, N1=14, N2=6):
        self.N1 = N1
        self.N2 = N2

    def __str__(self):
        name = type(self).__name__.lower()
        if self.N1 != 14 or self.N2 != 6:
            name += '_{}_{}'.format(self.N1, self.N2)
        return name
    
    def _load_internal(self, instrument, start_time, end_time, freq):
        adx = ADX(self.N1, self.N2)
        adxr = (adx + Ref(adx, self.N2)) / 2
        return adxr.load(instrument, start_time, end_time, freq)
